package com.karolak.math.formula;

import java.text.DecimalFormat;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;

import org.apache.commons.math3.complex.Complex;
import org.apache.commons.math3.transform.DftNormalization;
import org.apache.commons.math3.transform.FastFourierTransformer;
import org.apache.commons.math3.transform.TransformType;
import org.apache.commons.math3.util.ArithmeticUtils;

import com.karolak.math.MathUtils;

public class Fft {

	private double[] sample;
	private int harmonicCount;
	private Complex[] spectrum;
	private Map<Integer, double[]> harmonics = new LinkedHashMap<Integer, double[]>();
	private List<Entry<Integer, Complex>> significantFrequencies = new ArrayList<Entry<Integer, Complex>>();
	private double[] filtered;

	public void reset() {
		setSample(null);
		setHarmonicCount(0);
		setSpectrum(null);
		setFiltered(null);
		getHarmonics().clear();
		getSignificantFrequencies().clear();
	}

	public void calculate(double[] signal, int harmonicCount) {

		if (signal == null || signal.length == 0 || harmonicCount == 0) {
			return;
		}

		signal = getSignalPowerOfTwoLenth(signal);
		Complex[] spectrum = fft(signal);
		spectrum = normalizeSpectrum(spectrum);
		harmonicCount = harmonicCount <= spectrum.length ? harmonicCount : spectrum.length;

		setSample(signal);
		setSpectrum(spectrum);
		setHarmonicCount(harmonicCount);

		List<Entry<Integer, Complex>> sortedFrequencies = getSortedByMagnitude(getSpectrum());
		getSignificantFrequencies().addAll(sortedFrequencies.subList(0, harmonicCount));

		double[][] harmonics = new double[getHarmonicCount()][];
		for (int ix = 0; ix < getSignificantFrequencies().size(); ix++) {
			Entry<Integer, Complex> entry = getSignificantFrequencies().get(ix);
			double[] harmonic = calcHarmonic(ix);
			getHarmonics().put(entry.getKey(), harmonic);
			harmonics[ix] = harmonic;
		}
		setFiltered(MathUtils.add(harmonics));
	}

	// === implementation === //
	protected Complex[] fft(double[] sample) {
		FastFourierTransformer fft = new FastFourierTransformer(DftNormalization.STANDARD);
		return fft.transform(sample, TransformType.FORWARD);
	}

	protected Complex[] inFft(Complex[] input) {
		FastFourierTransformer fft = new FastFourierTransformer(DftNormalization.STANDARD);
		return fft.transform(input, TransformType.INVERSE);
	}

	public Complex[] normalizeSpectrum(Complex[] spectrum) {

		if (spectrum.length < 2) {
			return spectrum;
		}

		Complex[] spec = new Complex[spectrum.length];
		spec[0] = spectrum[0];
		for (int i = 1; i < spectrum.length / 2; i++) {
			spec[i] = spectrum[i].multiply(2);
		}
		Arrays.fill(spec, spectrum.length / 2, spectrum.length, Complex.ZERO);
		return spec;
	}

	public double[] calcHarmonic(int ix) {
		return MathUtils.toReal(inFft(calcSpectrum(ix)));
	}

	public Complex[] calcSpectrum(int ix) {
		Entry<Integer, Complex> entry = getSignificantFrequencies().get(ix);
		return calcSpectrum(getSpectrum().length, entry.getKey(), entry.getValue());
	}

	public Complex[] calcSpectrum(int size, int frequency, Complex magnitude) {
		Complex[] spectrum = new Complex[size];
		Arrays.fill(spectrum, Complex.ZERO);
		spectrum[frequency] = magnitude;
		return spectrum;
	}

	public double getLast(double x) {
		if (getFiltered() == null || getFiltered().length == 0) {
			return x;
		}
		return getFiltered()[getFiltered().length - 1];
	}
	
	public double getHrLast(Integer ix) {
		if (ix > getSignificantFrequencies().size() - 1) {
			return 0;
		}
		Entry<Integer, Complex> entry = getSignificantFrequencies().get(ix);
		double[] hr = getHarmonics().get(entry.getKey());
		return hr[hr.length - 1];
	}
	
	public double getFrequency(Integer ix) {
		double value = 0;
		if (ix > getSignificantFrequencies().size() - 1) {
			return value;
		}
		Entry<Integer, Complex> entry = getSignificantFrequencies().get(ix);
		value = entry.getKey();
		return value;
	}

	public double getPeriod(Integer ix) {
		double value = 0;
		if (ix > getSignificantFrequencies().size() - 1) {
			return 0;
		}
		Entry<Integer, Complex> entry = getSignificantFrequencies().get(ix);
		value = entry.getKey() == 0 ? Double.POSITIVE_INFINITY : getSampleLength() / entry.getKey();
		return value;
	}

	public double getMagnitude(Integer ix) {
		double value = 0;

		if (ix > getSignificantFrequencies().size() - 1) {
			return 0;
		}

		Entry<Integer, Complex> entry = getSignificantFrequencies().get(ix);
		value = entry.getValue().abs();
		return value;
	}

	public double getAmplitude(Integer ix) {
		return getMagnitude(ix) / getSpectrum().length;
	}

	public int getSampleLength() {
		return getSample() == null ? 0 : getSample().length;
	}

	protected List<Entry<Integer, Complex>> getSortedByMagnitude(Complex[] spectrum) {
		Map<Integer, Complex> spectrumMap = new HashMap<Integer, Complex>();
		List<Entry<Integer, Complex>> spectrumList = new ArrayList<Entry<Integer, Complex>>();
		for (int i = 0; i < spectrum.length; i++) {
			spectrumMap.put(i, spectrum[i]);
		}
		for (Entry<Integer, Complex> entry : spectrumMap.entrySet()) {
			spectrumList.add(entry);
		}

		Comparator<Map.Entry<Integer, Complex>> comp = new Comparator<Map.Entry<Integer, Complex>>() {
			public int compare(Entry<Integer, Complex> e0, Entry<Integer, Complex> e1) {
				Double d0 = Double.valueOf(e0.getValue().abs());
				Double d1 = Double.valueOf(e1.getValue().abs());
				return -d0.compareTo(d1);
			}
		};
		Collections.sort(spectrumList, comp);
		return spectrumList;
	}

	protected double[] getSignalPowerOfTwoLenth(double[] x) {
		int size = x.length;
		int s2 = getPowerOfTwo(size);
		if (s2 < size) {
			x = Arrays.copyOfRange(x, size - s2, size);
		}
		return x;
	}

	protected int getPowerOfTwo(int x) {
		int v = x;
		if (ArithmeticUtils.isPowerOfTwo(x)) {
			return x;
		}
		for (int i = x; i > 0; i--) {
			if (ArithmeticUtils.isPowerOfTwo(i)) {
				return i;
			}
		}
		return v;
	}

	// === get/set === //
	public double[] getFiltered() {
		return filtered;
	}

	protected void setFiltered(double[] filtered) {
		this.filtered = filtered;
	}

	protected double[] getSample() {
		return sample;
	}

	protected void setSample(double[] sample) {
		this.sample = sample;
	}

	public int getHarmonicCount() {
		return harmonicCount;
	}

	protected void setHarmonicCount(int harmonicCount) {
		this.harmonicCount = harmonicCount;
	}

	public Complex[] getSpectrum() {
		return spectrum;
	}

	protected void setSpectrum(Complex[] spectrum) {
		this.spectrum = spectrum;
	}

	public Map<Integer, double[]> getHarmonics() {
		return harmonics;
	}

	protected void setHarmonics(Map<Integer, double[]> harmonics) {
		this.harmonics = harmonics;
	}

	public List<Entry<Integer, Complex>> getSignificantFrequencies() {
		return significantFrequencies;
	}

	protected void setSignificantFrequencies(List<Entry<Integer, Complex>> significantFrequencies) {
		this.significantFrequencies = significantFrequencies;
	}

	public String describeHarmonic(int ix) {
		DecimalFormat df = new DecimalFormat("0.0000");
		StringBuilder sb = new StringBuilder();
		sb.append("H.").append(ix).append("{");
		sb.append("fr:").append(getFrequency(ix));
		sb.append(", per:").append(getPeriod(ix));
		sb.append(", mag:").append(df.format(getMagnitude(ix)));
		sb.append(", ampl:").append(df.format(getAmplitude(ix)));
		sb.append("}");
		return sb.toString();
	}
}
